import cv2
import os
import numpy as np
import pickle
import glob


def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


label_name = ["airplane",
              "automobile",
              "bird",
              "cat",
              "deer",
              "dog",
              "frog",
              "horse",
              "ship",
              "truck"]

train_file_list = glob.glob("E:\pythonProject\dataset\cifar10\cifar-10-batches-py\data_batch_*")
test_file_list = glob.glob("E:\pythonProject\dataset\cifar10\cifar-10-batches-py\\test_batch*")
# save_path = "E:\pythonProject\dataset\cifar10\cifar-10-batches-py\TRAIN"
save_path = "E:\pythonProject\dataset\cifar10\cifar-10-batches-py\TEST"
# 便利训练集文件
for file in test_file_list:
    # 解析文件dict
    file_dict = unpickle(file)

    # 便利文件中的每一个图片
    for im_index, im_data in enumerate(file_dict[b'data']):

        # 拿到对应标签
        im_label = file_dict[b'labels'][im_index]
        # 翻译标签
        im_label_name = label_name[im_label]

        # 拿到对应文件名
        im_filename = file_dict[b'filenames'][im_index]

        # 对图片构造修改RBG格式
        im_data = np.reshape(im_data, [3, 32, 32])
        im_data = np.transpose(im_data, [1, 2, 0])

        # cv2.imshow("im_data",cv2.resize(im_data,(200,200)))
        # cv2.waitKey(0)

        # 按标签名存储文件
        save_path_label = "{}/{}".format(save_path, im_label_name)
        if not os.path.exists(save_path_label):
            os.mkdir(save_path_label)
        # 存储图片
        cv2.imwrite("{}/{}".format(save_path_label, im_filename.decode("utf-8")), im_data)
